#ifndef __MF_H__
#define __MF_H__

#include <algorithm>
#include <iostream>
#include <string>
#include <vector>

#include "filter.h"

#include "Image.H"
#include "IO.H"
#include "Util.H"

using namespace pip;
using namespace std;

/// MEDIAN FILTERING \\\

/**
 * Performs median filtering on an image, using mirrored borders.
 *
 * @param image Image to filter
 * @param n Size of the mask (kernel)
 * @param squareMask Whether the mask is a square shape or a plus shape
 */
void medianFilter(Image<unsigned char> &image, const int &n, const bool &squareMask = true)
{
  const int kernelCentre(n / 2);
	Image<unsigned char> tempImage(image.bordered(kernelCentre, border_mirror));

  // Vector used to collect pixels over which median will be computed
  vector<unsigned char> pixels;

  for (int y(0); y < image.dim(1); ++y) {
    for (int x(0); x < image.dim(0); ++x) {
      for (int dy(-kernelCentre); dy <= kernelCentre; ++dy) {
        for (int dx(-kernelCentre); dx <= kernelCentre; ++dx) {
          // If we use a square mask, all pixels need to be processed.
          // If we use a plus mask, only pixels in the centre need to be processed.
          if (squareMask || dy == 0 || dx == 0) {
            pixels.push_back(tempImage(x + dx, y + dy));
          }
        }
      }

      // Iterator pointing to the median element of the vector.
	    vector<unsigned char>::iterator medianIterator(pixels.begin() + pixels.size() / 2);

      nth_element(pixels.begin(), medianIterator, pixels.end());
	    image(x, y) = *medianIterator;

      // Clear the vector, but keep its capacity
      pixels.erase(pixels.begin(), pixels.begin() + pixels.size());
    }
  }
}

/**
 * Performs fast median filtering on an image, using mirrored borders. The implementation is based
 * on the algorithm in Sonka, page 129-130. Adapted by Hugo Kuijf for extra speed.
 * It works for square kernels only, for plus-shaped kernels the above medianFilter method is called.
 * The same principle can be applied to plus-kernels, although the gain would not be that big as a
 * larger part of the kernel changes during the shift.
 *
 * Timings:
 * medianFilter as above        : 0.506 seconds
 * medianFilterFast as in Sonka : 0.032 seconds
 * medianFilterFast as below    : 0.027 seconds
 *
 * @param image Image to filter
 * @param n Size of the mask (kernel)
 * @param squareMask Whether the mask is a square shape or a plus shape
 */
void medianFilterFast(Image <unsigned char> &image, const int &n, const bool &squareMask = true)
{
  if (!squareMask) {
    return medianFilter(image, n, squareMask);
  }

  // Initialize variables
  int dx, dy, x, y, kernelCentre(n / 2), nm, initialNm, t(n * n / 2);
  Image<unsigned char> tempImage(image.bordered(kernelCentre, border_mirror));
  unsigned char m, initialM, pg;

  // Create initial histogram at (0,0) and compute the median
  int histogram[256], initialHistogram[256] = {0};

  //   Fill histogram
  for (dy = -kernelCentre; dy <= kernelCentre; ++dy) {
    for (dx = -kernelCentre; dx <= kernelCentre; ++dx) {
      initialHistogram[tempImage(dx, dy)]++;
    }
  }

  //   Compute initial mean (m) and number of pixels before the mean (nm)
  for (initialM = 0, initialNm = initialHistogram[initialM]; initialNm <= t; initialNm += initialHistogram[++initialM]);

  // Process whole image
  for (y = 0; y < image.dim(1); ++y) {
    // Copy the initial histogram, mean and nm for this row to the current histogram, mean and nm
    memcpy(histogram, initialHistogram, 256 * sizeof(int));
    m = initialM; nm = initialNm;

    // Shift histogram down. Only not at the first row, then we already have the correct histogram.
    if (y != 0) {
      // Create new initial histogram for this row, based on the initial histogram of the previous row
      for (dx = -kernelCentre; dx <= kernelCentre; ++dx) {
        // Clear top column
        pg = tempImage(dx, (y - 1) - kernelCentre);
        histogram[pg]--;
        if (pg <= m) --nm;

        // Fill bottom column
        pg = tempImage(dx, y + kernelCentre);
        histogram[pg]++;
        if (pg <= m) ++nm;
      }

      // Find the median. The median is at the location where the cumulative histogram (stored in nm) is equal to or larger than t.
      for (; nm <= t; ++m, nm += histogram[m]);
      for (; nm - histogram[m] > t; nm -= histogram[m], --m);

      // Copy the current histogram for this row, so the next row can use it as the initial histogram
      memcpy(initialHistogram, histogram, 256 * sizeof(int));
      initialM = m; initialNm = nm;
    }

    // Set the median
    image(0, y) = m;

    for (x = 1; x < image.dim(0); ++x) {
      for (dy = -kernelCentre; dy <= kernelCentre; ++dy) {
        // Clear left column
        pg = tempImage((x - 1) - kernelCentre, y + dy);
        histogram[pg]--;
        if (pg <= m) --nm;

        // Fill right column
        pg = tempImage(x + kernelCentre, y + dy);
        histogram[pg]++;
        if (pg <= m) ++nm;
      }

      // Find the median. The median is at the location where the cumulative histogram (stored in nm) is equal to t.
      for (; nm <= t; ++m, nm += histogram[m]);
      for (; nm - histogram[m] > t; nm -= histogram[m], --m);

      // Set the median
      image(x, y) = m;
    }
  }
}

/// SHARPENING \\\

/**
 * Initializes the Laplacian kernel
 *
 * @param kernel The Laplacian kernel
 * @param scale
 */
void initLaplacianKernel(Image<float> &kernel, const float &scale)
{
  kernel.resize(Dimension(3, 3));

  kernel(0,0) = -scale;
  kernel(1,0) = -scale;
  kernel(2,0) = -scale;

  kernel(0,1) = -scale;
  kernel(1,1) = 1 + 8 * scale;
  kernel(2,1) = -scale;

  kernel(0,2) = -scale;
  kernel(1,2) = -scale;
  kernel(2,2) = -scale;
}

/**
 * Sharpens the image by subtracting a Laplacian with factor C.
 * See also 5.35 in Sonka
 *
 * @param image Image to sharpen
 * @param C Factor
 */
void sharpen(Image<float> &image, const float &C)
{
  // Initialize the filter
  Image<float> kernel;
  initLaplacianKernel(kernel, C);

  Image<float> out;
  filter(image, out, kernel);

  image = out;
}

/// UNSHARP MASKING \\\

/**
 * Performs unsharp masking on an image.
 * 
 * @param image Image to filter
 * @param sigma
 * @param a
 */
void unsharpMasking(Image<float> &image, const float &sigma, const float &a)
{
  Image<float> gaussImage;
  gauss(image, gaussImage, sigma, 0, 0, 5.0);

  for (int y = 0; y < image.dim(1); ++y) {
    for (int x = 0; x < image.dim(0); ++x) {
      image(x, y) = (1.0f - a) * gaussImage(x, y) + a * image(x, y);
    }
  }
}

/// DENOISE \\\

/**
 * Generates a random number in the range [min, max]
 *
 * @param min Lower range value
 * @param max Upper range value
 * @return Random number
 */
float randomNumber(const float &min, const float &max)
{
  return min + ((max - min) * rand() / RAND_MAX);
}

/**
 * Generate a random number from a Gaussian distribution N(mean, sigma)
 *
 * @param mean
 * @param sigma
 * @return Random number
 */
float randomGaussian(const float &mean, const float &sigma)
{
  float U1, U2, V1, V2, S;

  do {
    U1 = randomNumber(0, 1);
    U2 = randomNumber(0, 1);
    V1 = 2 * U1 - 1;
    V2 = 2 * U2 - 1;
    S = V1 * V1 + V2 * V2;
  } while (S >= 1);

  float X = sqrt(-2 * log(S) / S) * V1;

  return (mean + sigma * X);
}

/**
 * Degrade an image with additive Gaussian noise.
 *
 * @param image Image to degrade
 * @param sigma
 */
void degradeGauss(Image<unsigned char> &image, const float &sigma)
{
  for (int y = 0; y < image.dim(1); ++y) {
    for (int x = 0; x < image.dim(0); ++x) {
      float randGauss = randomGaussian(0, sigma);

      int sign = randGauss < 0 ? -1 : 1,
        greyVal = (int)floor(abs(randGauss) + 0.5),
        noisyValue = image(x, y) + sign * greyVal;

      if (noisyValue < 0) {
        image(x, y) = 0;
      } else if (noisyValue > 255) {
        image(x, y) = 255;
      } else {
        image(x, y) = (unsigned char) noisyValue;
      }
    }
  }
}

/**
 * Degrade an image with impulse noise. Fraction denotes the
 * fraction of pixels affected by corruption.
 *
 * @param image Image to degrade
 * @param fraction
 */
void degradeImpulse(Image<unsigned char> &image, const float &fraction)
{
  for (int y = 0; y < image.dim(1); ++y) {
    for (int x = 0; x < image.dim(0); ++x) {
      float randUniform = randomNumber(-1, 1);

      if (abs(randUniform) < fraction) {
        image(x, y) = randUniform < 0 ? 0 : 255;
      }
    }
  }
}

/**
 * Computes the average squared difference between two input images.
 *
 * @param input First input image
 * @param secondInput Second input image
 * @return Average squared difference
 */
float averageSquaredDifference(const Image<unsigned char> &input, const Image<unsigned char> &secondInput)
{
  float diff(0);
  for (int y = 0; y < input.dim(1); ++y) {
    for (int x = 0; x < input.dim(0); ++x) {
      diff += (input(x,y) - secondInput(x,y)) * (input(x,y) - secondInput(x,y));
    }
  }

  return (diff / (input.dim(0) * input.dim(1)));
}

/**
 * Degrades an image with noise (additive Gaussian or Impulse noise),
 * and restores it again (by Gaussian blurring or median filtering).
 *
 * @param image Input image
 * @param degradedImage Output degraded image
 * @param degradedFilteredImage The 'restored' degraded output image
 * @param s1 Parameter for degrading (sigma for Gaussian noise or fraction for impulse noise)
 * @param s2 Parameter for restoring (sigma for blurring or n for median filtering)
 * @param filterMethod True for Gaussian blurring, false for median filtering
 */
void denoise(const Image<unsigned char> &image,
             Image<unsigned char> &degradedImage, 
             Image<unsigned char> &degradedFilteredImage, 
             const float &s1, 
             const float &s2, 
             const bool &filterMethod,
             void (*additiveNoise)(Image<unsigned char>&, const float &))
{
  // Make sure to check if filterMethod has an allowed value
  // Make sure that s1 is in a valid range, depending on the noise method used
  if(additiveNoise == degradeGauss && s1 <= 0) {
    cerr << "Incorrect value for s1, sigma must be > 0" << endl;
    return;
  } else if (additiveNoise == degradeImpulse && (s1 < 0 || s1 > 1)) {
    cerr << "Incorrect value for s1, sigma must be > 0 and < 1" << endl;
    return;
  }

  // Make sure s2 has valid value, it should be an odd number for median filtering
  if(!filterMethod && (int)s2 % 2 == 0) { // median filtering with even kernel size
    cerr << "Median filtering only allowed with odd kernel size" << endl;
    return;
  }

  additiveNoise(degradedImage, s1);

  if(filterMethod) { // Gaussian averaging
    Image<float> gaussImage, tempImage;

    convertUnsignedCharToFloat(degradedImage, tempImage);
    gauss(tempImage, gaussImage, s2);
    convertFloatToUnsignedChar(gaussImage, degradedFilteredImage);
  } else { // Median filtering
    degradedFilteredImage = degradedImage;
    medianFilter(degradedFilteredImage, (int)s2);
  }
}

/// CANNY EDGE DETECTION \\\

/**
 * Gradient magnitude ...
 *
 * @param in Input image
 * @param dx Output of X-derivative
 * @param dy Output of Y-derivative
 * @param gm Output of Gradient Magnitude
 * @param sigma
 */
void gradientMagnitude(const Image<unsigned char> &in, Image<float> &dx, Image<float> &dy, Image<float> &gm, const float &sigma)
{
  Image<float> image;
  gm.resize(in.dimension());

  convertUnsignedCharToFloat(in, image);

  gauss(image, dx, sigma, 1);
  gauss(image, dy, sigma, 0, 1);

  for(int y = 0; y < gm.dim(1); ++y) {
    for(int x = 0; x < gm.dim(0); ++x) {
      gm(x, y) = sqrt(dx(x, y) * dx(x, y) + dy(x, y) * dy(x, y));
    }
  }
}

/**
 * Canny edge detector ...
 *
 * @param in Input image
 * @param out Output of gradient magnitude for pixels that survive non-maximal suppression
 * @param sigma Scale for the gradient magnitude
 */
void cannyEdgeDetector(const Image<unsigned char> &in, Image<float> &out, const float &sigma)
{
  out.resize(in.dimension());
  out.fill(0);

  Image<float> dxImage, dyImage;
  Image<float> gmImage(in.dimension());

  // Compute gradient magnitude and edge direction
  gradientMagnitude(in, dxImage, dyImage, gmImage, sigma);

  // Apply non-maximum suppression in direction of the edge
  for(int y = 0; y < in.dim(1); ++y) {
    for(int x = 0; x < in.dim(0); ++x) {
      float gMag(gmImage(x, y));

      //if(gMag > 1.0E-10) { //if smaller, set to 0
        // calculate normal to edge according to formula 4.63 in Sonka
        // normal = derivative/gradientmagnitude
        float dx(dxImage(x, y) / gMag),
          dy(dyImage(x, y) / gMag);

        // do non-maximal suppression --> get only the strongest edge
        // Get gradientmagnitude interpolation value in positive and negative direction normal to edge
        // to check if these edges are stronger!
        float gMag1(linearInterpolation(gmImage, x + dx, y + dy)),
          gMag2(linearInterpolation(gmImage, x - dx, y - dy));

        if(gMag1 < gMag && gMag2 < gMag) { // x, y is maximum
          out(x, y) = gMag;
       // }
      } else {
        out(x, y) = 0;
      }
    }
  }
}

/**
 * Number of edge voxels ...
 *
 * @param filteredImage ...
 * @return ...
 */
int nrOfEdgeVoxels(const Image<float> &filteredImage)
{
  int sum(0);

  for(int y = 0; y < filteredImage.dim(1); ++y) {
    for(int x = 0; x < filteredImage.dim(0); ++x) {
      if(filteredImage(x, y) != 0) {
        ++sum;
      }
    }
  }

  return sum;
}


/// BODY OF THE PROGRAM \\\
/// You do not need to edit anything below this line!!
int applyMedian (int argc, char **argv)
{
	if (argc < 8) {
		error("", "Usage: MoreFiltering -operation Median -i <input> -o <output> "
		"-width N (and optional) -kernel S(quare) or P(lus)");
	}

	// Store commandline args in stl container
	vector<string> args(argv, argv + argc);

	// Variables to be read from the commandline
	string program, inputfile, outputfile, kernel("S");
	int width(3);

	// Read the arguments.
	getProgramName(args, program);
	getScalarArg(args, "-width", width);
	getScalarArg(args, "-i", inputfile);
	getScalarArg(args, "-o", outputfile);
	getScalarArg(args, "-kernel", kernel);

	bool square(true);
  if(kernel == "P") {
	   square = false;
  }

	cout << "Applying Median Filtering on " << inputfile << " using " << (square ? "square" : "plus-shaped") << " kernel with size " << width << "." <<  endl;
	Image<unsigned char> image;
	if(!importFile(image, inputfile)) {
		cerr << "Error loading " << inputfile << endl;
		return 1;
	}

	//getting min and max for assignment
	unsigned char min, max;
	getMinMaxImage(image, min, max);
	cout << "min of image before filtering: "<< (int) min << endl;
	cout << "max of image before filtering: "<< (int) max << endl;

	medianFilter(image, width, square);

	getMinMaxImage(image, min, max);
	cout << "min of image after filtering: "<< (int) min << endl;
	cout << "max of image after filtering: "<< (int) max << endl;

	if(!exportFile(image, outputfile)) {
		cerr << "Error writing " << outputfile << endl;
		return 1;
	}

	cout << "Successfully saved output as " << outputfile << "." << endl;

	return 0;
}

int applyUnsharp(int argc, char **argv)
{
	if (argc < 5) {
		error("", "Usage: MoreFiltering -operation UnsharpMask -i <input> -o <output> [-sigma s] [-A a]");
	}

	// Store commandline args in stl container
	vector<string> args(argv, argv + argc);

	// Variables to be read from the commandline
	string inputfile, outputfile, program;
	float sigma(2),
    a(1.5);

	// Read the arguments.
	getProgramName(args, program);
  getScalarArg(args, "-i", inputfile);
  getScalarArg(args, "-o", outputfile);
  getScalarArg(args, "-sigma", sigma);
  getScalarArg(args, "-A", a);

  cout << "Applying Unsharp Masking on " << inputfile << " using sigma: " << sigma << " and A: " << a << "." <<  endl;

  //Image<float> image;
  Image<unsigned char> imageUC;
  if (!importFile(imageUC, inputfile)) {
    cerr << "Error loading " << inputfile << "!" << endl;
    error(program, "Error reading image");
  }

  Image<float> image(imageUC.dimension());
	convertUnsignedCharToFloat(imageUC, image);

	//getting min and max for assignment
	float min, max;
	getMinMaxImage(image, min, max);
	cout << "min of image before filtering: " << min << endl;
	cout << "max of image before filtering: " << max << endl;

	unsharpMasking(image, sigma, a);

	getMinMaxImage(image, min, max);
	cout << "min of image after filtering: " << min << endl;
	cout << "max of image after filtering: " << max << endl;

  convertFloatToUnsignedChar(image, imageUC);

  if (!exportFile(imageUC, outputfile)) {
    error(program, "Error writing file");
  }

  cout << "Successfully saved output as " << outputfile << "." << endl;
  return 0;
}


int applySharpening(int argc, char **argv)
{
  if (argc < 6) {
    error("", "Usage: MoreFiltering -operation Sharpening -i <input> -o <output> [-factor F]");
  }

  // Store commandline args in stl container
  vector<string> args(argv, argv+argc);

  // Variables to be read from the commandline
  string inputfile, outputfile;
  string program;
  float C(0.7f);

	// Read the arguments.
  getProgramName(args, program);
  getScalarArg(args, "-i", inputfile);
  getScalarArg(args, "-o", outputfile);
  getScalarArg(args, "-factor", C);

  cout << "Applying Sharpening on " << inputfile << " using C = " << C << "." <<  endl;
  Image<unsigned char> imageUC;
  if (!importFile(imageUC, inputfile)) {
    cerr << "Error loading " << inputfile << "!" << endl;
    error(program, "Error reading image");
  }

	Image<float> image;
  image.resize(imageUC.dimension());

  convertUnsignedCharToFloat(imageUC, image);

	//getting min and max for assignment
	float min, max;
	getMinMaxImage(image, min, max);
	cout << "min of image before filtering: " << min << endl;
	cout << "max of image before filtering: " << max << endl;

  sharpen(image, C);

	getMinMaxImage(image, min, max);
	cout << "min of image after filtering: " << min << endl;
	cout << "max of image after filtering: " << max << endl;

  convertFloatToUnsignedChar(image, imageUC);

  if (!exportFile(imageUC, outputfile)) {
    error(program, "Error writing file");
  }

  cout << "Successfully saved output as " << outputfile << "." << endl;
  return 0;
}

int applyDenoise(int argc, char **argv)
{
  if (argc < 10) {
    error("", "Usage: MoreFiltering -operation Denoise -i <input> -o <degradedoutput> -o2 <blurredoutput> "
      "[-sigma1 s1] [-sigma2 s2] [-f a(veraging)/m(edian)] [-n g(auss)/i(mpulse)]");
  }

  // Store commandline args in stl container
  vector<string> args(argv, argv+argc);

  // Variables to be read from the commandline
	string inputfile, outputfile, outputfileblurred, noise, f, program;
  float s1(2),
    s2(2);

  // Read the arguments.
  getProgramName(args, program);
  getScalarArg(args, "-i", inputfile);
  getScalarArg(args, "-o", outputfile);
  getScalarArg(args, "-o2", outputfileblurred);
  getScalarArg(args, "-sigma1", s1);
  getScalarArg(args, "-sigma2", s2);
  getScalarArg(args, "-f", f);
  getScalarArg(args, "-n", noise);

  cout << "Applying Denoise on " << inputfile << " using sigma1: " << s1 << ", signma2: " << s2 << ", f: " << f << ", n" << noise  << "." <<  endl;

	bool filter(f == "a");

  Image<unsigned char> image;
  if (!importFile(image, inputfile)) {
    error(program, "Error reading image");
  }

  Image<unsigned char> degradedImage(image), degradedBlurredImage(image);

  if(noise == "g") {//gaussian noise
    denoise(image, degradedImage, degradedBlurredImage, s1,  s2, filter, degradeGauss);
  } else {
    denoise(image, degradedImage, degradedBlurredImage, s1,  s2, filter, degradeImpulse);
  }

  cout << "Average Squared Distance (ASD) between original and degraded image = " << averageSquaredDifference(image, degradedImage) << endl;
  cout << "Average Squared Distance (ASD) between original and restored image = " << averageSquaredDifference(image, degradedBlurredImage) << endl;

  if (!exportFile(degradedImage, outputfile)) {
    error(program, "Error writing file");
  }
	cout << "Successfully saved output as " << outputfile << "." << endl;

  if (!exportFile(degradedBlurredImage, outputfileblurred)) {
    error(program, "Error writing file");
  }
  cout << "Successfully saved output2 as " << outputfileblurred << "." << endl;

  return 0;
}

int applyCannyEdgeDection(int argc, char **argv)
{
  if (argc < 9) {
    error("", "Usage: MoreFiltering -operation CannyEdgeDetection -i <input> "
      "-o outputGM, -o2 outputCanny [-s scale]");
  }

  // Store commandline args in stl container
  vector<string> args(argv, argv+argc);

  // Variables to be read from the commandline
  string inputfile, outputfileGM, outputfileCanny, program;
  float scale(1.0f);

  // Read the arguments.
  getProgramName(args, program);
  getScalarArg(args, "-s", scale);
  getScalarArg(args, "-i", inputfile);
  getScalarArg(args, "-o", outputfileGM);
  getScalarArg(args, "-o2", outputfileCanny);

  cout << "Applying Canny Edge Detection on " << inputfile << " using scale: " << scale << "." <<  endl;

  Image<unsigned char> image;
  importFile(image, inputfile);

  cout << "Computing canny edges on scale " << scale << endl;
	Image<float> outCanny;
  cannyEdgeDetector(image, outCanny, scale);

  cout << "Number of edge voxels in canny output: " << nrOfEdgeVoxels(outCanny) << endl;

  Image<unsigned char> outCannyUC(outCanny.dimension());
  convertFloatToUnsignedChar(outCanny, outCannyUC);

  if(!exportFile(outCannyUC, outputfileCanny)) {
    error(program, "Error writing file");
  }

  // Calculate GM image
  Image<float> dxImage, 
    dyImage, 
    gmImage(image.dimension());

  // Compute gradient magnitude and edge direction
  gradientMagnitude(image, dxImage, dyImage, gmImage, scale);

  cerr << "exporting " << outputfileGM << endl;
  Image<unsigned char> outGMuc(gmImage.dimension());
  convertFloatToUnsignedChar(gmImage, outGMuc);
  if(!exportFile(outGMuc, outputfileGM)) {
    error(program, "Error writing file");
  }

  return 0;
}

#endif 
